# -*- coding: utf-8 -*-
# @Time    : 2019/12/16 8:32
# @Author  : zhouqiang
# @Site    : 
# @File    : code_processor.py
# @Software: PyCharm

"""
脚本说明: 通用个体处理工作需要的类抽象
"""

import random

import quant_researcher.quant.project_tool.common_var
from quant_researcher.quant.datasource_fetch.fund_api import fund_info
from quant_researcher.quant.datasource_fetch.index_api import index_info
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool import hammer
from quant_researcher.quant.project_tool import celebrity
from quant_researcher.quant.project_tool.db_operator import conn_dispatcher, my_mysql
from quant_researcher.quant.project_tool import assert_tool


class CodeProcessorV1(object):
    """
    通用个体处理工作需要的类抽象

    :param func_to_deal: 个体处理函数，应返回一个 list
    :param kwargs:
        - func_to_find，个体寻找函数，默认：一个提前定义好的函数（参考代码）
        - func_to_save，存储函数，如果不想对结果做任何处理，传入None，不传的话，会给一个默认值
            默认：一个提前定义好的，根据传入的参数直接写库的函数
        - kwds_to_find，个体寻找函数的关键字参数，默认：{}
        - kwds_to_deal，个体寻找函数的关键字参数，默认：{}
        - kwds_to_save，个体寻找函数的关键字参数，默认：{}
        - test_codes，list，用来测试的个体，如果非空，则不进行个体寻找，直接使用这里的代码进行计算
            默认：None
        - test_count，int，个体寻找完成后随机抽样的个数，默认：None
        - process_number，int，多进程处理时的进程数，默认：None（之后取配置）
        - batch_number，int，将找到的所有个体分批的批数，默认：进程数
        - prepare_these_conn，list，为个体处理函数和结果处理函数增加的数据库 conn，默认：[]
        - ignore_err_of_a_code，bool，个体处理函数运行时发生错误是否忽略，默认：False
        - handler_frequency，str，个体处理函数返回的结果批（batch）还是单独（code）处理，默认：batch
    """
    def __init__(self, func_to_deal, **kwargs):
        self.processor = func_to_deal
        self.hunter = kwargs.pop('func_to_find', find_codes)
        self.res_handler = kwargs.pop('func_to_save', self.simple_res_handler)

        self.kwds_to_find = kwargs.pop('kwds_to_find', {})
        self.kwds_to_deal = kwargs.pop('kwds_to_deal', {})
        self.kwds_to_save = kwargs.pop('kwds_to_save', {})

        self.test_codes = kwargs.pop('test_codes', None)
        self.test_count = kwargs.pop('test_count', None)
        self.p_count = kwargs.pop('process_number', None)
        if self.p_count is None:
            self.p_count = quant_researcher.quant.project_tool.common_var.MANY_P
        self.b_count = kwargs.pop('batch_number', self.p_count)
        self.prepare_these_conn = kwargs.pop('prepare_these_conn', [])
        self.ignore_err_of_a_code = kwargs.pop('ignore_err_of_a_code', False)
        self.handler_frequency = kwargs.pop('handler_frequency', 'batch')
        self.data_4_api = kwargs.pop('data_4_api', False)
        assert_tool.value_in_list(
            x=self.handler_frequency, the_list=['batch', 'code'],
            name='二选一：batch（分批处理个体process结果）；code（单独处理）'
        )

    @staticmethod
    def simple_res_handler(data, db, table, p_key, **kwargs):
        """
        简单结果处理函数，将数据写入数据库

        :param list_OR_df data: 数据
        :param str db: 库名
        :param str table: 表名
        :param list p_key: 主键
        :param kwargs:
            - conn_dict，连接字典
            - replace_nan，bool，是否自动替换数据中的 np.nan 为 None，默认：True
            - create_missing_cols，bool，是否自动替换对齐数据中的键，默认：False
        :return:
        """
        LOG.info(f'开始存储(行={len(data)})，DB={db}，表={table}，主键={p_key}')
        conn_dict = kwargs.pop('conn_dict')
        replace_nan = kwargs.pop('replace_nan', True)
        create_missing_cols = kwargs.pop('create_missing_cols', False)
        my_mysql.on_duplicate_key_update(
            a_lst=data, primary_key=p_key, table=table, conn=conn_dict[db],
            replace_nan=replace_nan, create_missing_cols=create_missing_cols
        )

    def assembled_processor(self, a_code, **kwargs):
        tmp_kwds = kwargs.copy()  # 浅拷贝，对多个参数的处理可能出现莫名问题
        tmp_kwds.update(self.kwds_to_deal)  # dict，尽量使用 update 来更新数据
        res = self.processor(a_code, **tmp_kwds)

        if not res:
            return

        if self.data_4_api:
            return res
        else:
            if self.res_handler is not None:
                tmp_kwds = kwargs.copy()
                tmp_kwds.update(self.kwds_to_save)
                return self.res_handler(res, **tmp_kwds)

    def batch_processor(self, a_batch, **kwargs):
        # 因为这个函数是在pickle之后运行，因此可以申请conn这种对象
        tmp_kwds = kwargs.copy()  # 浅拷贝，对多个参数的处理可能出现莫名问题
        tmp_kwds.update(self.kwds_to_deal)  # dict，尽量使用 update 来更新数据
        conn_dict = None
        if self.prepare_these_conn:
            conn_dict = conn_dispatcher.dispatcher(
                db_names=self.prepare_these_conn
            )
            tmp_kwds.update({'conn_dict': conn_dict})
        LOG.info(f'批处理单体(个数：{len(a_batch)})时，'
                 f'需要准备的DB连接：{self.prepare_these_conn}。\n'
                 f'传给单体处理函数的关键字：{celebrity.str_a_dict(tmp_kwds)}')
        batch_res = []
        try:
            for a_code in a_batch:
                try:
                    code_res = self.processor(a_code, **tmp_kwds)
                    if code_res:  # 单体处理函数可能正常结束并返回 “False” 值
                        assert_tool.value_is_instance(x=code_res, o_type=list)
                        batch_res.extend(code_res)
                except Exception as err:
                    if not self.ignore_err_of_a_code:
                        raise err
                    else:
                        LOG.warning(f'忽略在处理{a_code}发生的错误：{err}')
            LOG.info(f'批处理完后，'
                     f'关键字-1：{celebrity.str_a_dict(kwargs)}，'
                     f'关键字-2：{celebrity.str_a_dict(tmp_kwds)}')
            LOG.info(f'批处理结果长={len(batch_res)}，'
                     f'处理函数：{self.res_handler.__name__}')
            if self.res_handler is not None and batch_res:
                tmp_kwds = kwargs.copy()
                LOG.info(f'批处理完后，处理完数据库连接需求前，'
                         f'结果处理函数的关键字：{celebrity.str_a_dict(tmp_kwds)}')
                if self.prepare_these_conn:
                    tmp_kwds.update({'conn_dict': conn_dict})
                tmp_kwds.update(self.kwds_to_save)
                LOG.info(f'批处理完后，处理完数据库连接需求后，'
                         f'结果处理函数的关键字：{celebrity.str_a_dict(tmp_kwds)}')
                return self.res_handler(batch_res, **tmp_kwds)
        finally:
            if self.prepare_these_conn:
                conn_dispatcher.close_them(conn_dict)

    def run(self):
        LOG.warning(f'👯 WELCOME，进程数={self.p_count}。'
                    f'个体结果处理频率={self.handler_frequency}。'
                    f'批数={self.b_count}(注：此参数只在批处理时才有用)。'
                    f'测试代码={self.test_codes}，测试个数={self.test_count}。')
        if self.p_count > self.b_count:
            LOG.warning(f'进程数{self.p_count}比批数{self.b_count}大。'
                        '看来你的设置是没有get到这class的设计理念。'
                        '增加“批数设置”的目的是灵活调整每次存到数据库的行数：'
                        '如果数据量不大，则采用默认值，与进程数相同即可；'
                        '如果单体处理的行数很大，那么有必要缩减每次写库的行数，'
                        '此时需要增加批数。'
                        'ZQ注 on 2019年12月17日09:15:13')

        if self.test_codes is None:
            m_c = self.hunter(**self.kwds_to_find)
        else:
            m_c = self.test_codes

        if self.test_count is not None:
            if len(m_c) > self.test_count:
                m_c = random.sample(m_c, self.test_count)
            else:
                LOG.warning(f'尝试随机选取个数：{self.test_count}，'
                            f'却只找到需计算的个体数：{len(m_c)}')

        if self.handler_frequency == 'code':
            info = {'func': self.assembled_processor}
            if self.prepare_these_conn:
                info.update({'prepare_these_conn': self.prepare_these_conn})
            if not self.data_4_api:
                hammer.split_a_list_into_batches_and_deal_each_with_a_process(
                    the_list=m_c,
                    info=info,
                    how_many=self.p_count,
                    split_the_lst=False,
                    what_version='v2'
                )
            else:
                api_res = hammer.split_a_list_into_batches_and_deal_each_with_a_process(
                    the_list=m_c,
                    info=info,
                    how_many=self.p_count,
                    split_the_lst=False,
                    what_version='v2'
                )
                return api_res
        else:
            the_lst = hammer.split_a_list_into_batches(
                a_list=m_c, how_many=self.b_count
            )
            hammer.batch_process(
                a_list=the_lst,
                func=self.batch_processor,
                multiple_threads=False,
                pool_size=self.p_count
            )
        LOG.info('👮 ur done ...')


class CodesProcessorV2(object):
    def __init__(self, how_to_deal_with_a_code, how_to_handle_the_res_of_a_code, **kwargs):
        self.kwds_4_code_finder = kwargs.pop('kwds_4_code_finder', {})
        self.kwds_4_code_processor = kwargs.pop('kwds_4_code_processor', {})
        self.kwds_4_code_truck = kwargs.pop('kwds_4_code_truck', {})
        self.test_codes = kwargs.pop('test_codes', None)
        self.test_count = kwargs.pop('test_count', None)
        self.p_count = kwargs.pop('p_count', quant_researcher.quant.project_tool.common_var.MANY_P)
        self.set_basic_data_conn = kwargs.pop('set_basic_data_conn', True)
        self.set_derivative_data_conn = kwargs.pop(
            'set_derivative_data_conn', True
        )
        self.set_derivative_data_write_conn = kwargs.pop(
            'set_derivative_data_write_conn', True
        )
        self.how_to_find = kwargs.pop('how_to_find', find_codes)
        self.code_processor = how_to_deal_with_a_code
        self.code_truck = how_to_handle_the_res_of_a_code

    def code_process_core(self, a_code, **kwargs):
        tmp_kwds = kwargs.copy()  # 浅拷贝，对多个参数的处理可能出现莫名问题
        tmp_kwds.update(self.kwds_4_code_processor)  # dict，尽量使用 update 来更新数据
        res = self.code_processor(a_code, **tmp_kwds)
        if self.code_truck is not None:
            tmp_kwds = kwargs.copy()
            tmp_kwds.update(self.kwds_4_code_truck)
            return self.code_truck(res, **tmp_kwds)

    def run(self):
        if self.test_codes is None:
            m_c = self.how_to_find(**self.kwds_4_code_finder)
        else:
            m_c = self.test_codes

        if self.test_count is not None:
            if len(m_c) > self.test_count:
                m_c = random.sample(m_c, self.test_count)
            else:
                LOG.warning(f'尝试随机选取个数：{self.test_count}，'
                            f'却只找到需计算的个体数：{len(m_c)}')

        info = {
            'func': self.code_process_core,
            'set_basic_data_conn': self.set_basic_data_conn,
            'set_derivative_data_conn': self.set_derivative_data_conn,
            'set_derivative_data_write_conn': self.set_derivative_data_write_conn,
        }
        hammer.split_a_list_into_batches_and_deal_each_with_a_process(
            the_list=m_c, how_many=self.p_count, info=info
        )


def find_codes(code_class):
    if code_class == 'fund':
        all_codes = fund_info.get_fund_sname()
        all_codes = list(all_codes['fund_code'].values)
    elif code_class in ['stockindex', 'bondindex', 'stock_index', 'bond_index']:
        all_codes = index_info.get_all_index_code()
    elif code_class == 'manager':
        all_codes = fund_info.get_all_manager_code()
    else:
        raise NotImplementedError
    return all_codes
